import argparse
import torch
from transformers import AutoModelForImageClassification
from torchvision import transforms
from PIL import Image

# 设置命令行参数解析器
parser = argparse.ArgumentParser(description="Handwritten Digit Recognition")
parser.add_argument("image_path", type=str, help="Path to the input image")
args = parser.parse_args()

# 加载预训练模型
model_name = "farleyknight-org-username/vit-base-mnist"  # 使用指定的模型
model = AutoModelForImageClassification.from_pretrained(model_name)

# 定义图像转换
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小为 224x224
    transforms.Grayscale(num_output_channels=3),  # 转为 3 通道图像
    transforms.ToTensor(),         # 转为张量
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

# 加载和处理图像
image_path = args.image_path  # 从命令行接收的图像路径
image = Image.open(image_path)
image = transform(image).unsqueeze(0)  # 添加一个维度以适应模型输入

# 设置模型为评估模式
model.eval()

# 不需要计算梯度
with torch.no_grad():
    outputs = model(image)

# 获取预测的类
predicted_class = torch.argmax(outputs.logits, dim=1).item()
print(f"Predicted digit: {predicted_class}")